import numpy as np


class PNP:
    @staticmethod
    def dlt(object_points: np.ndarray, image_points: np.ndarray):
        rows = object_points.shape[0]
        A = np.zeros((2 * rows, 12))

        for i in range(rows):
            point1 = np.append(object_points[i, :], 1)
            A[2 * i, :4] = point1
            A[2 * i, -4:] = -point1 * image_points[i, 0]
            A[2 * i + 1, 4:-4] = point1
            A[2 * i + 1, -4:] = -point1 * image_points[i, 1]

        row2 = A.shape[0]
        col2 = A.shape[1]
        B = A.copy()
        for j in range(8, col2):
            B[:, j] = B[:, j] / np.max(B[:, j])

        temp = B @ np.linalg.pinv(A)
        A = temp @ A

        u, s, vh = np.linalg.svd(A)
        result = vh[-1, :]

        R1 = np.array([
            result[0: 3],
            result[4: 7],
            result[8: 11]
        ])
        u1, s1, vh1 = np.linalg.svd(R1)
        R = u1 @ vh1

        mean = np.mean(s1)

        t1 = np.array([result[3], result[7], result[11]])
        t = t1 / mean

        T_result = np.vstack((R.T, t)).T

        point1 = T_result @ np.append(object_points[0, :], 1)
        if point1[2] < 0:
            T_result = -T_result

        return T_result
    @staticmethod
    def dlt2(object_points: np.ndarray, image_points: np.ndarray, camera_matrix):
        rows = object_points.shape[0]
        inv_camera_matrix = np.linalg.inv(camera_matrix)
        A = np.zeros((2 * rows, 12))

        temp = np.zeros((rows, 3))
        for i in range(rows):
            point2 = np.append(image_points[i, :], 1)
            temp[i, :] = inv_camera_matrix @ point2

        for i in range(rows):
            point1 = np.append(object_points[i, :], 1)
            A[2 * i, :4] = point1
            A[2 * i, -4:] = -point1 * temp[i, 0]
            A[2 * i + 1, 4:-4] = point1
            A[2 * i + 1, -4:] = -point1 * temp[i, 1]

        u, s, vh = np.linalg.svd(A)
        result = vh[-1, :]

        R1 = np.array([
            result[0: 3],
            result[4: 7],
            result[8: 11]
        ])
        u1, s1, vh1 = np.linalg.svd(R1)
        R = u1 @ vh1

        mean = np.mean(s1)

        t1 = np.array([result[3], result[7], result[11]])
        t = t1 / mean

        T_result = np.vstack((R.T, t)).T

        point1 = T_result @ np.append(object_points[0, :], 1)
        if point1[2] < 0:
            T_result = -T_result

        T_result = np.vstack((T_result, [0, 0, 0, 1]))
        return T_result
